#include <fstream>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <deque>
#include <string>
#include <iostream>


#define DEBUG_LEVEL 6
#include <debug_func.h>

#include <blast_results.h>
#include <bcpp_read_seq.h>
#include <portable_timer.h>
#include <codeml.h>
#include <codeml_results.h>
#include <progress_indicator.h>
#define _MY_STL_MACROS
#include <stlhelper.h>
#include <ultoa.h>
#include <auto_restore_dir.h>
#include <print_error.h>
#include <seq_pairs.h>

#include <boost/lexical_cast.hpp>
#include <boost/array.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>

#include "get_arg.h"


using std::vector;
using std::deque;
using std::ofstream;
using std::ifstream;
using std::istream;
using std::istringstream;
using std::string;
using std::runtime_error;
using std::cerr;
using std::cin;
using std::cout;
using BioCPP::t_collect_sequences;
using BioCPP::eSEQ_NODUPLICATES;
using BioCPP::eSEQ_DUPLICATES;
using BioCPP::t_collect_sequences_report_errors_to_strm;
using boost::filesystem::current_path;
using boost::filesystem::path;
using boost::lexical_cast;
using bcpp_align::t_align_results;

struct t_align_results_ids
{
	unsigned seq_id1;
	unsigned seq_id2;
	t_align_results blast_result;
	t_align_results_ids(unsigned n1, unsigned n2, const t_align_results& blast_result_)
			:seq_id1(n1), seq_id2(n2), blast_result(blast_result_){}
};


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	bitscore

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

/*
//________________________________________________________________________________________

//	get_seq_bitscore

//		calculate blosum raw bitscores for one stretch of a sequence pair

//________________________________________________________________________________________
// helper struct
struct op_calc_bitscore
{
	char	(*gap_penalty)[256];
	char	(*lookup_table)[256];
	int	bitscore;
	char prev_a;
	char prev_b;
	op_calc_bitscore(char gap_penalty_[256][256], char lookup_table_[256][256]):
			gap_penalty(gap_penalty_), lookup_table(lookup_table_),
			bitscore(0), prev_a('-'), prev_b('-'){}
	void operator()(char a, char b)
	{
		// lookup bitscore
		bitscore += lookup_table[a][b];

		// gap opening penalty
		bitscore += gap_penalty[prev_a][a];
		bitscore += gap_penalty[prev_b][b];
		// save current char to calculate gap opening penalty
		prev_a = a;
		prev_b = b;

	}
};

int get_seq_bitscore(const std::string& seq1, const std::string& seq2,
							unsigned beg, unsigned end, unsigned& len)
{
	if (!seq1.length())
		return 0;
	// skip past initial '-'
	beg = seq1.find_first_not_of('-', beg);
	beg = seq2.find_first_not_of('-', beg);

	// skip past final '-'
	end = seq1.find_last_not_of('-',  end);
	end = seq2.find_last_not_of('-',  end);
	if (beg == string::npos ||
		end == string::npos || end < beg)
		return 0;
	assert(beg >= 0);
	assert(end < seq1.length());
	assert(seq2.length() == seq1.length());
	len += end + 1 - beg;
	return (for_each_pair(	seq1.begin() + beg, seq1.begin() + end + 1,
							seq2.begin() + beg,
							op_calc_bitscore(gap_penalty, lookup_table))
			).bitscore;
}
//________________________________________________________________________________________

//	initialize_bitscore_lookup_table

//________________________________________________________________________________________
void initialize_bitscore_lookup_table(char lookup_table[256][256], char gap_penalty[256][256])
{
	// bitscore lookup table
	memset (lookup_table, 0, 256 * 256 * sizeof (char));
	for (unsigned i = 0; i < sizeof (blosum_chars)-1; ++i)
		for (unsigned j = 0; j < sizeof (blosum_chars)-1; ++j)
		lookup_table[blosum_chars[i]][blosum_chars[j]] = BLOSUM_MATRIX[i][j];

	// gap extension penalty
	for (unsigned i = 0; i < sizeof(blosum_chars)-1; ++i)
	{
		lookup_table[blosum_chars[i]]['-'] = -1;
		lookup_table['-'][blosum_chars[i]] = -1;
	}

	// gap initiation penalty
	memset (gap_penalty, 0, 256 * 256 * sizeof (char));
	for (unsigned i = 0; i < sizeof(blosum_chars)-1; ++i)
		gap_penalty[blosum_chars[i]]['-'] = -11;
}


#define BLOSUM80
#ifdef BLOSUM60
#define GAP_EXIST_PENALTY -11
#define BLOSUM_MATRIX blosum62
#else
#ifdef BLOSUM80
#define GAP_EXIST_PENALTY -10
#define BLOSUM_MATRIX blosum80
#endif
#endif
#include "blosum_data.h"

  */



//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	KaKs

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888









struct t_remove_stops
{
	static char table1 [256];
	static char table2 [256];
	static char table3 [256];
	static bool initialized;
	t_remove_stops()
	{
		if (!initialized)
			initialize();
	}
	void initialize()
	{
		std::fill(table1, table1 + 256, 2);
		std::fill(table2, table2 + 256, 1);
		std::fill(table3, table3 + 256, 2);
		table1[unsigned('t')] = 0;
		table2[unsigned('a')] = 0;
		table2[unsigned('g')] = 1;
		table3[unsigned('a')] = 0;
		table3[unsigned('g')] = 1;
		initialized = true;
	}

	unsigned validate(std::string& sequence)
	{
		unsigned cnt_stops = 0;
		assert(sequence.length() % 3 == 0);
		//throw std::runtime_error("cDNA sequence length not in multiple of 3.");
		string::iterator i = sequence.begin();
		string::iterator end = sequence.end();
		while (i != end)
		{
			if (table1[unsigned(*i)] + table2[unsigned(*(i+1))] + table3[unsigned(*(i+2))] < 2)
			{
				*i = *(i+1) = *(i+2) = 'N';
				cnt_stops++;
			}
			i += 3;
		}
		return cnt_stops;
	}
};

char t_remove_stops::table1 [256];
char t_remove_stops::table2 [256];
char t_remove_stops::table3 [256];
bool t_remove_stops::initialized = false;


//________________________________________________________________________________________

//		insert_gaps_into_sequence_pair

//________________________________________________________________________________________
bool insert_gaps_into_sequence_pair(	const t_align_results&		pair_data,
										string&						sequence1,
										string&						sequence2,
										string&						error_str)
{
	debug_func_cerr(5);
	try
	{

		if (!pair_data.align_len)
		{
			error_str = "Zero length alignment!";
			return false;
		}


		if (!sequence1.length() || !sequence2.length())
		{
			error_str = "Sequences missing.";
			return false;
		}

		if (!expand_gaps_cdna(pair_data.gaps1, sequence1) ||
			!expand_gaps_cdna(pair_data.gaps2, sequence2))
		{
			error_str = "Alignment gaps do not match sequence.";
			return false;
		}
		if (sequence1.length() != sequence2.length())
		{
			error_str = "Sequence lengths mismatch after inserting gaps. Lengths = "
							+ ultoa(sequence1.length()) + ", " + ultoa(sequence2.length());
			return false;
		}
		if (pair_data.align_beg * 3 > sequence1.length())
		{
			error_str = "Align start too large.";
			return false;
		}

		// remove non-aligned region
		sequence1.erase(0, pair_data.align_beg * 3);
		sequence2.erase(0, pair_data.align_beg * 3);

		if (pair_data.align_len * 3 > sequence1.length())
		{
			error_str = "Align length too large (alignment length = " +
							ultoa(pair_data.align_len)+ ", sequence length = " +
							ultoa(sequence1.length()) + ")";
			return false;
		}

		sequence1.erase(pair_data.align_len * 3);
		sequence2.erase(pair_data.align_len * 3);

		t_remove_stops().validate(sequence1);
		t_remove_stops().validate(sequence2);
		return true;

	}
	catch (...)
	{
		error_str = "Unknown gap insertion error\n";
		return false;
	}

}




//________________________________________________________________________________________

//		print_on_error

//			for debugging

//________________________________________________________________________________________
void print_on_error(t_program_args& args,
					const string& orig_sequence1,
					const string& orig_sequence2,
					const string& sequence1,
					const string& sequence2,
					const string& gap_str1, const string& gap_str2,
					const string& name1,
					const string& name2,
					unsigned align_len)
{
	debug_func_cerr(5);
	cerr << "names1 = " << name1.length() << "\t" << name1 << "\n";
	cerr << "names2 = " << name2.length() << "\t" << name2 << "\n";
	cerr << "orig_sequence1 = " << orig_sequence1.length() << "\t" << orig_sequence1 << "\n";
	cerr << "orig_sequence2 = " << orig_sequence2.length() << "\t" << orig_sequence2 << "\n";
	cerr << "gap_str1 = " << gap_str1 << "\n";
	cerr << "gap_str2 = " << gap_str2 << "\n";
	cerr << "sequence1 = " << sequence1.length() << "\t" << sequence1 << "\n";
	cerr << "sequence2 = " << sequence2.length() << "\t" << sequence2 << "\n";
	cerr << "align len = " << align_len * 3 << "\n\n";
	cerr.flush();
}






//________________________________________________________________________________________

//		get_kaks_results

//________________________________________________________________________________________
unsigned get_kaks_results(	t_collect_sequences&				seq_db,
							const deque<t_align_results_ids>&	alignment_data,
							t_progress_indicator& 				dots,
							t_program_args&						args,
							unsigned&							cnt_hung_paml_restarted)
{
	debug_func_cerr(5);
	unsigned cnt_failures = 0;

	vector<t_codeml_results> results;

	deque<t_align_results_ids>::const_iterator i = alignment_data.begin();
	deque<t_align_results_ids>::const_iterator end = alignment_data.end();
	for (; i != end; ++i, ++dots)
	{
		string error_str;
		string seq1 = seq_db.sequence(i->seq_id1);
		string seq2 = seq_db.sequence(i->seq_id2);
		if (args.aligned ||
			insert_gaps_into_sequence_pair(i->blast_result, seq1, seq2, error_str))
		{
			results.clear();
			if (BioCPP::kaks_analyse(seq1, seq2, 0, 0, results, args.paml + " kaks.ctl", 
									 args.yn00 ? BioCPP::eYN00 : BioCPP::eCODEML, 
									 cnt_hung_paml_restarted))
			{			
				cout	<<	seq_db.name(i->seq_id1)	<<	'\t'
						<<	seq_db.name(i->seq_id2)	<<	'\t'
						<<	results[0]				<<	'\t'
						<<	i->blast_result			<<	'\n';
				continue;
			}
			else
			{
				error_str = args.yn00 ? "yn00" : "codeml";
				error_str += " failed to complete its analysis within 5 minutes 3 times "
							 "running for this sequence pair";
			}
		}

		// first failure
		if (args.results_header && !cnt_failures)
			args.KAKS_LOG()
					<< seq_pairs::		column_names() << "\t"
					<< t_align_results::column_names() << "\t"
					<<	"Error\n"
					<< seq_pairs::		column_name_underscores() << "\t"
					<< t_align_results::column_name_underscores() << "\t"
					<<	"-----\n";

		++cnt_failures;
		//
		//	Error:
		//
		// save to db
		args.KAKS_LOG()
				<<	seq_db.name(i->seq_id1)	<<	'\t'
				<<	seq_db.name(i->seq_id2)	<<	'\t'
				<<	i->blast_result			<<	'\t'
				<<	error_str				<<	'\n';
		dots.print_indicator('!');


	}
	return alignment_data.size() - cnt_failures;
}


//________________________________________________________________________________________

//		do_parse_pairs

//________________________________________________________________________________________
struct do_parse_pairs
{
	const t_collect_sequences&		seq_db;
	deque<t_align_results_ids>& 	alignment_data;
	t_program_args&					args;
	std::string						comment_letters;
	do_parse_pairs(	const t_collect_sequences&			seq_db_,
					deque<t_align_results_ids>& 		alignment_data_,
					t_program_args&						args_,
					const std::string&					comment_letters_):

			seq_db(seq_db_), 
			alignment_data(alignment_data_),
			args(args_),
			comment_letters(comment_letters_){}
	void print_error(unsigned cnt_line, const string& prot_name) const
	{
		std_print_error(args.CERR(), "The sequence for \"" + prot_name +
								"\" on line #" + ultoa(cnt_line) +
								" could not be found");
	}
	void operator()(istream& is, unsigned cnt_line)
	{
		// get data for one protein pair
		string name1, name2;

		// comment if false
		string comment_line;
		if (!seq_pairs::parse_ids(is, name1, name2, comment_line, comment_letters,
								  seq_pairs::ePARSE_IDS_ONLY))
			return;

		t_align_results aligndata;
		is >> aligndata;

		if (!is)
			throw std::runtime_error("The alignment data is not in the correct format");
		
		// if names are already in sequence file, must be indexed
		unsigned index1, index2;
		if (!seq_db.get_index_of_name(name1,  index1))
		{
			print_error(cnt_line, name1);
			return;
		}
		if (!seq_db.get_index_of_name(name2,  index2))
		{
			print_error(cnt_line, name2);
			return;
		}

		alignment_data.push_back(t_align_results_ids(index1, index2, aligndata));
	}
};





void check_seq_same_len(t_collect_sequences& seq_db, unsigned i, unsigned j)
{
	if (seq_db.sequence(i).length() != seq_db.sequence(j).length())
		throw runtime_error("All sequences in a multiple alignment should be of "
							"the same length. The "
							"sequence #" + ultoa(i+1) + " [" + seq_db.name(i)
							+ "] has " + ultoa(seq_db.sequence(i).length())
							+ " residues, while "
							"sequence #" + ultoa(j+1) + " [" + seq_db.name(j) +
							+ "] has " + ultoa(seq_db.sequence(j).length())
							+ " residues.");

}



void create_empty_alignment_data_for_seq_pairs(	t_collect_sequences& seq_db,
												deque<t_align_results_ids>& alignment_data,
												t_program_args& args)
{
	//
	// check number of sequences
	//
	if (!seq_db.size())
		return;
	if (seq_db.size() == 1)
		throw runtime_error("Only 1 cDNA sequence was specified.\n");
	if (!args.all_vs_all && seq_db.size() % 2 != 0)
		throw runtime_error("Pairwise comparisons require an even number of sequences.\n"
							+ ultoa(seq_db.size()) + " sequences were supplied.\n");

	//
	//	check multiple of 3
	//
	for (unsigned i = 0; i < seq_db.size() ; ++i)
	{
		unsigned len = seq_db.sequence(i).length();
		if (len % 3 != 0)
			throw runtime_error("Aligned cDNA sequences must have integral number of "
								"codons.\nSequence " + seq_db.name(i) + " has " +
								ultoa(len) + " bases.\n");
	}

	//
	// create fake blast results
	//
	t_align_results empty_blast_result;
	if (args.all_vs_all)
	{
		// check all same length
		for (unsigned i = 1; i < seq_db.size(); ++i)
			check_seq_same_len(seq_db, 0, i);

		// construct all against all
		for (unsigned i = 1; i < seq_db.size(); ++i)
			for (unsigned j = 0; j < i; ++j)
			{
				get_align_results_from_cdna_sequence_pairs(empty_blast_result,
												seq_db.sequence(i),
												seq_db.sequence(j));
				alignment_data.push_back(t_align_results_ids(i, j,  empty_blast_result));

			}
		args.VERBOSE() << "\t" << alignment_data.size() << " all vs all sequence pairs.\n";
	}
	else
	{
		for (unsigned i = 0; i < seq_db.size(); i += 2)
		{
			check_seq_same_len(seq_db, i, i + 1);
			get_align_results_from_cdna_sequence_pairs(empty_blast_result,
											seq_db.sequence(i),
											seq_db.sequence(i + 1));

			// found both names in sequence files
			alignment_data.push_back(t_align_results_ids(i, i + 1,  empty_blast_result));
		}
		args.VERBOSE() << "\t" << alignment_data.size() << " successive sequence pairs.\n";
	}
}


void read_sequences(t_collect_sequences& seq_db, t_program_args& args)
{
	{
		ifstream f_cdna_data(args.file_cdna.c_str());
		if (!f_cdna_data)
			throw std::runtime_error("Could not open " + args.file_cdna +
										" to retrieve cDNA sequences");

		args.VERBOSE() << "\tReading cDNA sequences from [ " << args.file_cdna << " ]...\n";
		t_progress_indicator dots(args.VERBOSE(), 10000);
		t_collect_sequences_report_errors_to_strm report_errors(args.CERR());
		string valid_chars = "ACGTMRWSYKVHDBN";
		if (args.aligned)
			valid_chars += '-';
		seq_db.init_from_stream(f_cdna_data, args.regex, dots, report_errors, valid_chars, true);
	}
	args.VERBOSE() << "\t" << seq_db.size() << " sequences\n";
}

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		Main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main(int argc, char* argv[])
{
	t_program_args args;

	try
	{

		debug_func_cerr(5);

		// don't use C streams
		std::ios::sync_with_stdio(false);


		// get command line arguments
		if (!process_command_line_options(argc, argv, args))
			return 1;


		//
		// read sequences
		//
		t_collect_sequences seq_db(args.aligned ? eSEQ_DUPLICATES : eSEQ_NODUPLICATES);
		read_sequences(seq_db, args);
		if (!seq_db.size())
		{
			std_print_warning(args.CERR(), "No sequences were supplied.\n");
			return 0;
		}


		args.VERBOSE(2)	<< "\tUsing PAML from [ " << args.paml << " ].\n";





		//
		//	Run Kaks in temp directory
		//
		// use kaks_tmp as temp directory if not specified
		string tmp_file_path = args.tempdir.length() ?
										args.tempdir :
										current_path().string() + "/kaks_tmp";





		//
		// read sequences alignment data
		//
		deque<t_align_results_ids> alignment_data;
		if (args.aligned)
			create_empty_alignment_data_for_seq_pairs(seq_db, alignment_data, args);
		else
		{
			//
			// read sequences alignment data
			//
			args.VERBOSE() << "\tReading alignment data from STDIN...\n";
			do_parse_pairs parse_pairs(seq_db,  alignment_data, args, args.comment_letters);
			seq_pairs::read_data(parse_pairs, "STDIN", args.VERBOSE());
		}





		bool remove_tmp_file = false;
		{
			//
			// save original directory for restore later
			//
			t_auto_restore_dir save_directory;
			prepare_temp_directory(tmp_file_path,
										false, 0.0, /*fixed_kappa, args.kappa,*/
										args.yn00 ? BioCPP::eYN00 : BioCPP::eCODEML, 
										args.f61  ? BioCPP::eF61 : BioCPP::eF3x4,
										remove_tmp_file);
			args.VERBOSE(2) << "\tUsing [ " + tmp_file_path +
										" ] as the temp directory for KaKs analysis.\n";




			//
			// Write header
			//
			if (args.results_header)
				cout
						<< seq_pairs::			column_names() << "\t"
						<< t_codeml_results::	column_names() << "\t"
						<< t_align_results::	column_names() << "\n"
						<< seq_pairs::			column_name_underscores() << "\t"
						<< t_codeml_results::	column_name_underscores() << "\t"
						<< t_align_results::	column_name_underscores() << "\n";





			//
			// Calculate kaks
			//
			unsigned cnt_kaks = 0;
			unsigned cnt_hung_paml_restarted = 0;
			{
				t_progress_indicator dots(args.VERBOSE(), 1, alignment_data.size());
				dots.use_timer();
				cnt_kaks = get_kaks_results(seq_db, alignment_data, dots, args, cnt_hung_paml_restarted);
			}
			args.VERBOSE() << "\tKa/Ks calculated for " << cnt_kaks	<< " sequence pairs.\n";

			if (cnt_hung_paml_restarted)
			{
				args.CERR() << "\t" << cnt_hung_paml_restarted
							<< "\tPAML jobs hung and had to be restarted.\n";
			}
		}


		//
		// Removing temp directory
		//
		if (remove_tmp_file && !args.keep_tempdir)
		{
			args.VERBOSE(2) << "\tRemoving temporary directory [ " << tmp_file_path	<< " ].\n";
			boost::filesystem::remove_all(tmp_file_path);
		}


	}
	catch (std::exception& e)
	{
		std_print_error(args.CERR(), e.what());
		return 1;
	}

    cout << "#dNdS EOF\n";
	return 0;
}

